/**
 * @author Gao Jian
 * @email gaoj@njupt.edu.cn
 * @create date 2021-06-03 10:38:05
 * @desc 连接破碎的小块
 */

#include "defs.h"

__global__ void filter_IS_i(unsigned char *data, unsigned int nw,
                            unsigned int nh, unsigned int nd) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    unsigned char dat[5];
    dat[0] = data[idx];
    dat[1] = data[idx + nz * 1];
    dat[2] = data[idx + nz * 2];
    dat[3] = data[idx + nz * 3];
    dat[4] = data[idx + nz * 4];

    if (dat[1] + dat[2] + dat[3] + dat[4] <= 2 && dat[0] == 1)
      dat[0] = 0;

    if (dat[0] + dat[1] + dat[2] + dat[3] >= 2 && dat[4] == 0)
      dat[4] = 1;

    if (dat[0] == 1 && dat[1] == 0)
      dat[1] = 1;
    if (dat[1] == 1 && dat[2] == 0)
      dat[2] = 1;
    if (dat[2] == 1 && dat[3] == 0)
      dat[3] = 1;
    // if (dat[3] == 1 && dat[4] == 0)
    //   dat[4] = 1;

    if (dat[4] == 0 && dat[3] == 1)
      dat[3] = 0;
    if (dat[3] == 0 && dat[2] == 1)
      dat[2] = 0;
    if (dat[2] == 0 && dat[1] == 1)
      dat[1] = 0;
    if (dat[1] == 0 && dat[0] == 1)
      dat[0] = 0;

    data[idx] = dat[0];
    data[idx + nz * 1] = dat[1];
    data[idx + nz * 2] = dat[2];
    data[idx + nz * 3] = dat[3];
    data[idx + nz * 4] = dat[4];
  }
}

void filter_IS(unsigned char *data, unsigned int nw, unsigned int nh,
               unsigned int nd) {
  unsigned nz = nw * nh * nd * sizeof(unsigned char);
  unsigned char *d_data = NULL;
  cudaMalloc((void **)&d_data, nz);
  cudaMemcpy(d_data, data, nz, cudaMemcpyHostToDevice);

  filter_IS_i<<<NUM_BLOCKS, NUM_THREADS>>>(d_data, nw, nh, nd);

  cudaMemcpy(data, d_data, nz, cudaMemcpyDeviceToHost);
  cudaFree(d_data);
}
